module evolve

  use dimensions

  complex, private      :: ExpKv(jv, jz)
  complex, private      :: ExpKx(jx, jDir)
  complex, private      :: ExpKy(jy, jDir)
  complex, private      :: ExpKz(jz)

  real, private         ::  weight 

contains

!-------------------------------------------------------------------------------------
subroutine evolve_init()
!-------------------------------------------------------------------------------------

  use parameters, only : dt, kz, vPhase, ux0, uy0, Dxxxx, Dzzzz, iunit, SideLoss

  use mpi_wrap

  implicit none

  integer              :: iDir, ikz, iv


  write(msg,*) char(10), ' Initializing evolve module'
  call mpi_wrap_msg

  !-- initialization --

  weight = 1./(jDir*npDir)

  ExpKz(:) = cexp( -iunit * kZvec(:) * (-vPhase) * dt )
 
  do iv=1,jv
     ExpKv(iv,:) = cexp( -iunit*kZvec(:)*vvec(iv)*dt - Dzzzz * kZvec(:)**4 *dt  - SideLoss*dt  )
  enddo


  do iDir = 1,jDir
     ExpKx(:,iDir) = cexp( - iunit * kXvec(:)*ux0(iDir) * dt  )
     ExpKy(:,iDir) = cexp( - iunit * kYvec(:)*uy0(iDir) * dt  )
     !ExpKx(:,iDir) = cexp( - iunit * kXvec(:)*ux0(iDir) * dt - Dxxxx * kXvec(:)**4 *dt )
     !ExpKy(:,iDir) = cexp( - iunit * kYvec(:)*uy0(iDir) * dt - Dxxxx * kYvec(:)**4 *dt )
  enddo

end subroutine evolve_init


!-------------------------------------------------------------------------------------
subroutine evolve_F_long_advect(dfu, halfstep)
!-------------------------------------------------------------------------------------

  use fft_3d
  use fft_wrap

  implicit none

  complex, intent(inout)   :: dfu(jv, jz, jy, jx, jDir)
  logical, intent(in)      :: halfstep
  integer                  :: ix, iy, iz, iv, iDir


 !------------------------------------------

  !-- advect in z: df/dt = -v df/dz + Dzzzz d4f/dz4 - SideLoss

  !call fft_forward_dfu_z(dfu)
  call dfftw_execute_dft(plan_z5forward, dfu, dfu)
  dfu = dfu/jz

  if (halfstep) then

     do iDir = 1,jDir
        do ix=1,jx
           do iy=1, jy
              !$OMP PARALLEL SHARED(dfu, ExpKv)
              !$OMP WORKSHARE
              dfu(:,:,iy,ix,iDir) = dfu(:,:,iy,ix,iDir)*sqrt(ExpKv)
              !$OMP END WORKSHARE NOWAIT
              !$OMP END PARALLEL
           enddo
        enddo
     enddo

  else


     do iDir = 1,jDir
        do ix=1,jx
           do iy=1,jy
              !$OMP PARALLEL SHARED(dfu, ExpKv)
              !$OMP WORKSHARE
              dfu(:,:,iy,ix,iDir) = dfu(:,:,iy,ix,iDir)*ExpKv
              !$OMP END WORKSHARE NOWAIT
              !$OMP END PARALLEL 
           end do
        enddo
     enddo
 

  endif

  !call fft_backward_dfu_z(dfu)
  call dfftw_execute_dft(plan_z5backward, dfu, dfu)

end subroutine evolve_F_long_advect


!-------------------------------------------------------------------------------------
subroutine evolve_F_trans_advect(dfu, uxMat, uyMat, f0vec)
!-------------------------------------------------------------------------------------

  use parameters, only : Dxxxx, dt, iBcPerpPeriodic, ux0, uy0, iLinVlasov

  implicit none

  complex, intent(inout)   :: dfu(jv, jz, jy, jx, jDir)
  complex, intent(in)      :: uxMat(jz, jy, jx, jDir)
  complex, intent(in)      :: uyMat(jz, jy, jx, jDir)
  real,    intent(in)      :: f0vec(jv)

  real                     :: axVar(jx+1), fxVec(jx+1)
  real                     :: ayVar(jy+1), fyVec(jy+1)
  integer                  :: ix, iy, iz, iv, iDir
  integer                  :: jjx, jjy, iBcFlag

 !------------------------------------------

  if (iLinVlasov) then
     call evolve_F_trans_advect_lin(dfu, uxMat, uyMat, f0vec)
     return
  endif


  if (iBcPerpPeriodic) then
     jjx = jx
     jjy = jy
     iBcFlag = 1
  else  
     jjx = jx + 1
     jjy = jy + 1
     iBcFlag = 0
  end if


  do iDir = 1,jDir

     !-- advect in x: d(f+f0)/dt = -d/dx [(f+f0) Ux(x,y)] --

     do iy=1,jy
        do iz=1,jz                      
           axVar(1:jx) = uxMat(iz,iy,:,iDir)           
           axVar(jx+1) = ux0(iDir)          
           do iv=1,jv
              fxVec(1:jx) = dfu(iv,iz,iy,:,iDir) + weight * f0vec(iv)
              fxVec(jx+1) = weight * f0vec(iv)
              call UpWind2ndOrder_Variable_a( fxVec,  dt, dx, axVar, jjx, &
                   weight*f0vec(iv),  weight*f0vec(iv), iBcFlag )
              dfu(iv,iz,iy,:,iDir) = fxVec(1:jx)  - weight * f0vec(iv) 
           enddo
        enddo
     end do

     !-- advect in y: d(f+f0)/dt = -d/dy [(f+f0) Uy(x,y)] --

     do ix=1,jx
        do iz=1,jz                      
           ayVar(1:jy) = uyMat(iz,:,ix,iDir)           
           ayVar(jy+1) = uy0(iDir)
           do iv=1,jv
              fyVec(1:jy) = dfu(iv,iz,:,ix,iDir) + weight * f0vec(iv)
              fyVec(jy+1) = weight * f0vec(iv)
              call UpWind2ndOrder_Variable_a( fyVec,  dt, dy, ayVar, jjy, &
                   weight*f0vec(iv),  weight*f0vec(iv), iBcFlag )
              dfu(iv,iz,:,ix,iDir) = fyVec(1:jy)  - weight * f0vec(iv) 
           enddo
        enddo
     end do


   enddo ! iDir

end subroutine evolve_F_trans_advect



!-------------------------------------------------------------------------------------
subroutine evolve_F_trans_advect_lin(dfu, uxMat, uyMat, f0vec)
!-------------------------------------------------------------------------------------

  use fft_3d
  use fft_wrap

  use parameters, only:  dt, iunit

  implicit none

  complex, intent(inout)   :: dfu(jv, jz, jy, jx, jDir)
  complex, intent(in)      :: uxMat(jz, jy, jx, jDir)
  complex, intent(in)      :: uyMat(jz, jy, jx, jDir)
  real,    intent(in)      :: f0vec(jv)


  complex    :: dummy1Dx(jx)
  complex    :: dummy1Dy(jy)

  integer    :: ix, iy, iz, iv, iDir

   
  !-- advect dfu in y:  d/dt (dfu) = - Uy0 d/dy(dfu) - f0 d/dy(Uy)
  !-- in the future, if we ever run this version, two terms need to be combined
  !-- and penscil ffts of Uy need to be replaced by block fft.

  !-- y-advection, first term: d/dt (dfu) = - Uy0 d/dy(dfu)

  !call fft_forward_dfu_y(dfu)
  call dfftw_execute_dft(plan_y5forward, dfu, dfu)
  dfu = dfu/jy

  do iDir = 1,jDir
      do ix=1,jx
        do iz=1,jz
           do iv=1,jv

             dfu(iv,iz,:,ix,iDir) = dfu(iv,iz,:,ix,iDir) * ExpKy(:,iDir)

           end do
        end do
     end do
  end do

  !call fft_backward_dfu_y(dfu)
  call dfftw_execute_dft(plan_y5backward, dfu, dfu)


 !-- x-advection, first term: d/dt (dfu) = - Ux0 d/dx(dfu)

  !call fft_forward_dfu_x(dfu)
  call dfftw_execute_dft(plan_x5forward, dfu, dfu)
  dfu = dfu/jx

  do iDir = 1,jDir
     do iy=1,jy
        do iz=1,jz
           do iv=1,jv

              dfu(iv,iz,iy,:,iDir) =  dfu(iv,iz,iy,:,iDir) * ExpKx(:,iDir)

           end do
        end do
     end do
  end do

  !call fft_backward_dfu_x(dfu)
  call dfftw_execute_dft(plan_x5backward, dfu, dfu)


  !-- y-advection, second term:  d/dt (dfu) = - f0 d/dy(Uy)

  do iDir = 1,jDir
     do ix=1,jx
        do iz=1,jz

           dummy1Dy(:) = uyMat(iz,:,ix,iDir)
           call dfftw_execute_dft(plan_yforward, dummy1Dy, dummy1Dy)
           dummy1Dy=dummy1Dy/float(jy)

           dummy1Dy(:) = - iunit * kYvec(:) * dummy1Dy(:) * dt

           call dfftw_execute_dft(plan_ybackward, dummy1Dy, dummy1Dy)


           do iv=1,jv
              dfu(iv,iz,:,ix,iDir) =  dfu(iv,iz,:,ix,iDir) + dummy1Dy(:)*weight*f0vec(iv)
           end do

        end do
     enddo
  end do


  !-- x-advection, second term:  d/dt (dfu) = - f0 d/dx(Ux)


  do iDir = 1,jDir  
     do iy=1,jy
        do iz=1,jz

           dummy1Dx(:) = uxMat(iz,iy,:,iDir)
           call dfftw_execute_dft(plan_xforward, dummy1Dx, dummy1Dx)
           dummy1Dx=dummy1Dx/float(jx)

           dummy1Dx(:) = - iunit * kXvec(:) * dummy1Dx(:) * dt

           call dfftw_execute_dft(plan_xbackward, dummy1Dx, dummy1Dx)

           do iv=1,jv
              dfu(iv,iz,iy,:,iDir) =  dfu(iv,iz,iy,:,iDir) + dummy1Dx(:)*weight*f0vec(iv)
           end do

        end do
     enddo
  end do


end subroutine evolve_F_trans_advect_lin



!-------------------------------------------------------------------------------------
subroutine evolve_F_accel(dfu, Ez2D, f0vec)
!-------------------------------------------------------------------------------------

  use fft_wrap
  use fft_3d
  use parameters, only : dt, Dvvvv, CollisionRate, iunit

  implicit none

  complex, intent(inout)   :: dfu(jv, jz, jy, jx, jDir)
  complex, intent(in)      :: Ez2D(jz,jy,jx)
  real,    intent(in)      :: f0vec(jv)

  complex                  :: f0kV(jv), DvvvvTerm(jv), b0(jv), c0(jv)
  complex                  :: b(jv), c(jv), e(jv)
  integer                  :: ix, iy, iz, iv, iDir

  real                     :: rc

  rc = CollisionRate

  f0kV(:)=f0vec(:)*weight      
  call dfftw_execute_dft(plan_vforward, f0kV, f0kV)
  f0kV=f0kV/jv

  f0kV =  kVvec(:)*f0kV(:)

  DvvvvTerm(:) = Dvvvv*kVvec(:)**4
  b0(:)  =  DvvvvTerm(:) + rc*kVvec(:)**2

  !-- acceleration: df/dt = Ez df/dv

   !call fft_forward_dfu_v(dfu)
   call dfftw_execute_dft(plan_v5forward, dfu, dfu)
   dfu = dfu/float(jv)

   do iDir = 1,jDir
      do ix=1,jx
         do iy=1,jy

            !-- before applying hyperviscosity, subtract z-average, on Denis's demand
            c0(:) =  sum(dfu(:,:,iy,ix,iDir), 2)/jz
            c0(:) = DvvvvTerm(:)*c0(:)

            !$OMP PARALLEL PRIVATE(iz, b, e, c)
            !$OMP DO             
            do iz=1,jz
               b(:) =  -b0(:) - iunit*kVvec(:)*Ez2D(iz,iy,ix)
               e(:) =   cexp(b(:)*dt)
               c(:) =  -iunit*f0kV(:)*Ez2D(iz,iy,ix) + c0(:)
               c(:) =   c(:)/b(:)
               c(1) =   0
               dfu(:,iz,iy,ix,iDir) =  dfu(:,iz,iy,ix,iDir)*e + c * (e - 1)
            enddo
            !$OMP END DO
            !$OMP END PARALLEL

         enddo
      enddo
   enddo
  
   !call fft_backward_dfu_v(dfu)
   call dfftw_execute_dft(plan_v5backward, dfu, dfu)

end subroutine evolve_F_accel


!-------------------------------------------------------------------------------------
subroutine evolve_Uz(dfu, f0vec, uzMat)
!-------------------------------------------------------------------------------------

  implicit none

  complex, intent(in)      :: dfu(jv, jz, jy, jx, jDir)
  real,    intent(in)      :: f0vec(jv)
  real,    intent(out)     :: uzMat(jz, jy, jx, jDir)


  real       :: fVsVvec(jv)
  integer    :: ix, iy, iz, iv, iDir


  !-- integrate over "v" to find longtitudinal velocity Uz -- 
       
  do iDir = 1,jDir

     do ix = 1,jx
        do iy = 1,jy
           do iz = 1,jz 
              fVsVvec(:) = dfu(:,iz,iy,ix,iDir) + weight * f0vec(:)   ! weight or simpy 0.5?
              uzMat(iz,iy,ix,iDir) = sum( fVsVvec * vVec ) / sum( fVsVvec )
           enddo
        enddo
     enddo

  enddo
   

end subroutine evolve_Uz


!-------------------------------------------------------------------------------------
subroutine evolve_U_long_advect(uxMat, uyMat, uzMat)
!-------------------------------------------------------------------------------------

  use fft_3d

  use parameters, only : dt, iFourier, iWaveFrame


  implicit none

  complex, intent(inout)   :: uxMat(jz, jy, jx, jDir)
  complex, intent(inout)   :: uyMat(jz, jy, jx, jDir)
  real,    intent(in)      :: uzMat(jz, jy, jx, jDir)


  complex    :: dummy1Dz(jz)

  real       :: uxVecZ(jz)
  real       :: uyVecZ(jz)
  real       :: uzVecZ(jz)

  integer    :: ix, iy, iz, iv, iDir

  if( iFourier==1 .and. iWaveFrame==0) return
  
  !-- advect U in z:  dUx/dt = -Uz dUx/dz, dUy/dt = -Uz dUy/dz

  if( iFourier==1 ) then                   ! FT with linearization in uz

     !--  advect ux in z since <u_z> approximately = -vphase =  V0

     !call fft_forward_u_z(uxMat)
     call dfftw_execute_dft(plan_z4forward, uxMat, uxMat)
     uxMat = uxMat/jz

     !call fft_forward_u_z(uyMat)
     call dfftw_execute_dft(plan_z4forward, uxMat, uxMat)
     uyMat = uyMat/jz

     do iDir = 1,jDir
        do ix=1,jx
           do iy=1,jy
              uxMat(:,iy,ix,iDir) = uxMat(:,iy,ix,iDir) *  ExpKz(:)
              uyMat(:,iy,ix,iDir) = uyMat(:,iy,ix,iDir) *  ExpKz(:)
           end do
        end do
     end do

     !call fft_backward_u_z(uxMat)
     call dfftw_execute_dft(plan_z4backward, uxMat, uxMat)

     !call fft_backward_u_z(uyMat)
     call dfftw_execute_dft(plan_z4backward, uxMat, uxMat)


  else                                     ! upwind or central differences

     do iDir = 1,jDir
        do ix=1,jx
           do iy=1,jy

              uzVecZ(:) = uzMat(:,iy,ix,iDir)
              uxVecZ(:) = uxMat(:,iy,ix,iDir)
              uyVecZ(:) = uyMat(:,iy,ix,iDir)

              call SubUpwind_NonCon_Form( uxVecZ, dt, dz, uzVecZ, jz, &
                   0.0, 0.0, 1, -iFourier, dummy1Dz )

              call SubUpwind_NonCon_Form( uyVecZ, dt, dz, uzVecZ, jz, &
                   0.0, 0.0, 1, -iFourier, dummy1Dz )

              uxMat(:,iy,ix,iDir)= uxVecZ(:)
              uyMat(:,iy,ix,iDir)= uyVecZ(:)

           enddo
        enddo
     enddo


  end if

end subroutine evolve_U_long_advect

!-------------------------------------------------------------------------------------
subroutine evolve_U_trans_advect_lin(uxMat, uyMat)
!-------------------------------------------------------------------------------------

  use fft_3d
  use parameters, only : dt


  implicit none

  complex, intent(inout)   :: uxMat(jz, jy, jx, jDir)
  complex, intent(inout)   :: uyMat(jz, jy, jx, jDir)


  integer    :: ix, iy, iz, iv, iDir

   
  !-- advect Ux and Uy in y:  dUx/dt = -Uy0 dUx/dy,   dUy/dt = -Uy0 dUy/dy

  !call fft_forward_u_y(uxMat)
  call dfftw_execute_dft(plan_y4forward, uxMat, uxMat)
  uxMat = uxMat/jy

  !call fft_forward_u_y(uyMat)
  call dfftw_execute_dft(plan_y4forward, uyMat, uyMat)
  uyMat = uyMat/jy

  do iDir = 1,jDir
     do ix=1,jx
        do iz=1,jz

           uxMat(iz,:,ix,iDir) = uxMat(iz,:,ix,iDir) * ExpKy(:,iDir)
           uyMat(iz,:,ix,iDir) = uyMat(iz,:,ix,iDir) * ExpKy(:,iDir)

        end do
     enddo
  end do

  !call fft_backward_u_y(uxMat)
  call dfftw_execute_dft(plan_y4backward, uxMat, uxMat)

  !call fft_backward_u_y(uyMat)
  call dfftw_execute_dft(plan_y4backward, uyMat, uyMat)


  !-- advect Ux and Uy in x:  dUx/dt = -Ux0 dUx/dx,   dUy/dt = -Ux0 dUy/dx

  !call fft_forward_u_x(uxMat)
  call dfftw_execute_dft(plan_y4forward, uxMat, uxMat)
  uxMat = uxMat/jx

  !call fft_forward_u_x(uyMat)
  call dfftw_execute_dft(plan_y4forward, uyMat, uyMat)
  uyMat = uyMat/jx

  do iDir = 1,jDir   
     do iy=1,jy
        do iz=1,jz
           uxMat(iz,iy,:,iDir) = uxMat(iz,iy,:,iDir) *  ExpKx(:,iDir)
           uyMat(iz,iy,:,iDir) = uyMat(iz,iy,:,iDir) *  ExpKx(:,iDir)
        enddo
     enddo
  enddo

  !call fft_backward_u_x(uxMat)
  call dfftw_execute_dft(plan_x4backward, uxMat, uxMat)

  !call fft_backward_u_x(uyMat)
  call dfftw_execute_dft(plan_x4backward, uyMat, uyMat)


end subroutine evolve_U_trans_advect_lin


!-------------------------------------------------------------------------------------
subroutine evolve_U_trans_advect(uxMat, uyMat)
!-------------------------------------------------------------------------------------

  use parameters, only : ux0, uy0, dt, iBcPerpPeriodic, iLinHydro

  implicit none

  complex, intent(inout)   :: uxMat(jz, jy, jx, jDir)
  complex, intent(inout)   :: uyMat(jz, jy, jx, jDir)

  real       :: dummy1Dy(jy)
  real       :: dummy1Dx(jx)

  real       :: uxVecX(jx+1)
  real       :: uyVecX(jx+1)

  real       :: uxVecY(jy+1)
  real       :: uyVecY(jy+1)

  integer    :: ix, iy, iz, iv, iDir
  integer    :: jjx, jjy, iBcFlag

  if (iLinHydro) then
    call evolve_U_trans_advect_lin(uxMat, uyMat)
    return
  endif

  if (iBcPerpPeriodic) then
     jjx = jx
     jjy = jy
     iBcFlag = 1
  else  
     jjx = jx + 1
     jjy = jy + 1
     iBcFlag = 0
  end if

 

  !-- advect U in y:  dUx/dt = -Uy dUx/dy,   dUy/dt = -Uy dUy/dy


  do iDir = 1,jDir
   
     do ix=1,jx
        do iz=1,jz
             
           uxVecY(1:jy) = uxMat(iz,:,ix,iDir)
           uyVecY(1:jy) = uyMat(iz,:,ix,iDir)
           uxVecY(jy+1) = ux0(iDir)
           uyVecY(jy+1) = uy0(iDir)

           call SubUpwind_NonCon_Form( uxVecY, dt, dy, uyVecY, jjy, &
                ux0(iDir), ux0(iDir), iBcFlag, 1, dummy1Dy )

           call NonLinearAdvection( uyVecY, dt, dy, jjy, &
                uy0(iDir), uy0(iDir), iBcFlag )  

           uxMat(iz,:,ix,iDir) = uxVecY(1:jy)
           uyMat(iz,:,ix,iDir) = uyVecY(1:jy)

        end do
        enddo
        
  end do

 
  !-- advect U in x:  dUx/dt = -Ux dUx/dx,   dUy/dt = -Ux dUy/dx

  do iDir = 1,jDir

     do iy=1,jy 
        do iz=1,jz

           uxVecX(1:jx) = uxMat(iz,iy,:,iDir)
           uyVecX(1:jx) = uyMat(iz,iy,:,iDir)
           uxVecX(jx+1) = ux0(iDir)
           uyVecX(jx+1) = uy0(iDir)


           call SubUpwind_NonCon_Form( uyVecX, dt, dx, uxVecX, jjx, &
                uy0(iDir), uy0(iDir), iBcFlag, 1, dummy1Dx )

           call NonLinearAdvection( uxVecX, dt, dx, jjx, &
                ux0(iDir), ux0(iDir), iBcFlag )  
 
           uxMat(iz,iy,:,iDir) = uxVecX(1:jx)
           uyMat(iz,iy,:,iDir) = uyVecX(1:jx)

        end do
     end do

  end do

end subroutine evolve_U_trans_advect


!-------------------------------------------------------------------------------------
 subroutine evolve_U_accel(uxMat, uyMat, Ex2D, Ey2D)
!-------------------------------------------------------------------------------------

  use parameters, only : dt

  implicit none

  complex, intent(inout)   :: uxMat(jz, jy, jx, jDir)
  complex, intent(inout)   :: uyMat(jz, jy, jx, jDir)

  complex, intent(in)      :: Ex2D(jz,jy,jx)
  complex, intent(in)      :: Ey2D(jz,jy,jx)

  integer                  :: iDir

  !-- acceleration:  dUx/dt = Ex,  dUy/dt = Ey


  do iDir = 1,jDir
   
     uxMat(:,:,:,iDir) = uxMat(:,:,:,iDir) + Ex2D * dt
     uyMat(:,:,:,iDir) = uyMat(:,:,:,iDir) + Ey2D * dt
         
  enddo
 

end subroutine evolve_U_accel



!-------------------------------------------------------------------------------------

end module evolve
